#include <time.h>
#include <string>
#include <iostream>
#include <iomanip>
#include <fstream>
#include <stdlib.h>
#include <TTree.h>
#include <TChain.h>
#include "Plot2D.h"
using namespace std;

int make_plot(string inputlist="input_list.txt") {

    //<Get input list>
    ifstream filelist;
    filelist.open(inputlist.c_str());
    if(!filelist.good()) {
        cout<<"ERROR: Cannot find the input filelist, now quit!"<<endl;
        exit(-1);
    }

    //<Define root TChain, tree name comes from input>
    string treename[4]={"tree_incl_4e", "tree_incl_2e2mu", "tree_incl_2mu2e", "tree_incl_4mu"};
    TChain *chain_data[4], *chain_mc[4];
    for( int i=0; i<4; i++ ) {
       cout << "tree name: " << treename[i].c_str() << endl;
       chain_data[i] = new TChain(treename[i].c_str());
       chain_mc[i] = new TChain(treename[i].c_str());
    }

    //<Add trees into TChain>
    string file, tree;
    while(!filelist.eof()) {
        getline(filelist,file);
        if(file.size()==0) continue; //remove the blank lines
        cout<<"Add file \""<<file<<"\""<<endl;
        if(file.find("data")!=string::npos) {
		   for(int i=0; i<4; i++) { 
           tree=file+"/"+treename[i];
           cout << "data:" << tree.c_str() << endl;
           chain_data[i]->Add(tree.c_str()); 
           }
        }
        if(file.find("mc")!=string::npos) {
           for(int i=0; i<4; i++) { 
           tree=file+"/"+treename[i];
           cout << "mc:" << tree.c_str() << endl;
           chain_mc[i]->Add(tree.c_str()); 
           }
        }
    }

    //<Set up loading options to make it loop faster>
    for( int i=0; i<4; i++ ) {
       cout << "Entries data: " << chain_data[i]->GetEntries() << endl;
       cout << "Entries mc: " << chain_mc[i]->GetEntries() << endl;
       chain_data[i]->SetCacheSize(10000000);
       chain_mc[i]->SetCacheSize(10000000);
    }

    
    Plot2D *myplot = new Plot2D();
    for(int i=0; i<4; i++) {
       string stream="data";
       myplot->Init(chain_data[i]);
       long nevts = chain_data[i]->GetEntries();
       for(long j=0; j<nevts; j++) {
           myplot->LoadEntry(j);
           myplot->ProcessEntry(j, stream, i);
       }
    }
    for(int i=0; i<4; i++) {
       string stream="mc";
       myplot->Init(chain_mc[i]);
       long nevts = chain_mc[i]->GetEntries();
       for(long j=0; j<nevts; j++) {
           myplot->LoadEntry(j);
           myplot->ProcessEntry(j, stream, i);
       }
    }
 
    myplot->Terminate();

    //<delete the instance at the end>
    delete myplot;
 

    //<The END>
    return 0;

}

int main(int argc, char* argv[]){
    if(argc==1) make_plot();
    else if(argc==2) make_plot(argv[1]);

    return 0;
}
